package com.controller;

import com.service.MessageService;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.context.ApplicationContext;

import javax.websocket.RemoteEndpoint;
import javax.websocket.Session;
import java.io.IOException;
import java.lang.reflect.Field;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.*;

/**
 * ChatWebSocket单元测试
 */
@ExtendWith(MockitoExtension.class)
public class ChatWebSocketTest {

    @Mock
    private ApplicationContext applicationContext;

    @Mock
    private MessageService messageService;

    @Mock
    private Session session1;

    @Mock
    private Session session2;

    @Mock
    private RemoteEndpoint.Basic remoteEndpoint1;

    @Mock
    private RemoteEndpoint.Basic remoteEndpoint2;

    private ChatWebSocket chatWebSocket;

    @BeforeEach
    void setUp() {
        // 清理静态状态
        clearStaticState();
        
        chatWebSocket = new ChatWebSocket();
        chatWebSocket.setApplicationContext(applicationContext);
        
        // 设置基本的mock行为
        when(applicationContext.getBean(MessageService.class)).thenReturn(messageService);
    }

    @AfterEach
    void tearDown() {
        // 清理静态状态
        clearStaticState();
    }

    /**
     * 清理ChatWebSocket的静态状态
     */
    private void clearStaticState() {
        try {
            Field sessionMapField = ChatWebSocket.class.getDeclaredField("sessionMap");
            sessionMapField.setAccessible(true);
            java.util.Map<String, Session> sessionMap = (java.util.Map<String, Session>) sessionMapField.get(null);
            if (sessionMap != null) {
                sessionMap.clear();
            }
        } catch (Exception e) {
            // 忽略反射异常
        }
    }

    @Test
    void testOnOpen() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session1.isOpen()).thenReturn(true);
        
        // 测试用户连接
        chatWebSocket.onOpen(session1, "user1");
        
        // 验证在线用户数
        assertEquals(1, ChatWebSocket.getOnlineUserCount());
        assertTrue(ChatWebSocket.getOnlineUsers().contains("user1"));
        
        // 验证MessageService被获取
        verify(applicationContext).getBean(MessageService.class);
    }

    @Test
    void testOnMessage() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session2.getBasicRemote()).thenReturn(remoteEndpoint2);
        when(session1.isOpen()).thenReturn(true);
        when(session2.isOpen()).thenReturn(true);
        
        // 先建立连接
        chatWebSocket.onOpen(session1, "user1");
        chatWebSocket.onOpen(session2, "user2");
        
        // 发送消息
        String message = "Hello, WebSocket!";
        chatWebSocket.onMessage(message, "user1");
        
        // 验证消息被保存到数据库
        verify(messageService).saveMessage(message, "user1", "all");
        
        // 验证消息被广播给其他用户（user2应该收到消息）
        verify(remoteEndpoint2).sendText("用户 user1: " + message);
        // 发送者不应该收到自己的消息
        verify(remoteEndpoint1, never()).sendText(anyString());
    }

    @Test
    void testOnClose() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session2.getBasicRemote()).thenReturn(remoteEndpoint2);
        when(session1.isOpen()).thenReturn(true);
        when(session2.isOpen()).thenReturn(true);
        
        // 建立连接
        chatWebSocket.onOpen(session1, "user1");
        chatWebSocket.onOpen(session2, "user2");
        
        assertEquals(2, ChatWebSocket.getOnlineUserCount());
        
        // 关闭连接
        chatWebSocket.onClose("user1");
        
        // 验证用户被移除
        assertEquals(1, ChatWebSocket.getOnlineUserCount());
        assertFalse(ChatWebSocket.getOnlineUsers().contains("user1"));
        assertTrue(ChatWebSocket.getOnlineUsers().contains("user2"));
    }

    @Test
    void testOnError() {
        // 测试错误处理
        Throwable error = new RuntimeException("Test error");
        
        // 不应该抛出异常
        assertDoesNotThrow(() -> {
            chatWebSocket.onError(session1, error);
        });
    }

    @Test
    void testSendMessageToUser() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session1.isOpen()).thenReturn(true);
        
        // 建立连接
        chatWebSocket.onOpen(session1, "user1");
        
        // 发送消息给指定用户
        String message = "Private message";
        boolean result = ChatWebSocket.sendMessageToUser("user1", message);
        
        assertTrue(result);
        verify(remoteEndpoint1).sendText(message);
    }

    @Test
    void testSendMessageToNonExistentUser() {
        // 发送消息给不存在的用户
        boolean result = ChatWebSocket.sendMessageToUser("nonexistent", "message");
        
        assertFalse(result);
    }

    @Test
    void testSendMessageToClosedSession() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session1.isOpen()).thenReturn(true);
        
        // 建立连接
        chatWebSocket.onOpen(session1, "user1");
        
        // 模拟会话关闭
        when(session1.isOpen()).thenReturn(false);
        
        // 发送消息应该失败
        boolean result = ChatWebSocket.sendMessageToUser("user1", "message");
        
        assertFalse(result);
    }

    @Test
    void testGetOnlineUserCount() {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session2.getBasicRemote()).thenReturn(remoteEndpoint2);
        when(session1.isOpen()).thenReturn(true);
        when(session2.isOpen()).thenReturn(true);
        
        // 初始状态应该为0（由于@BeforeEach中的清理）
        assertEquals(0, ChatWebSocket.getOnlineUserCount());
        
        // 添加用户
        chatWebSocket.onOpen(session1, "user1");
        assertEquals(1, ChatWebSocket.getOnlineUserCount());
        
        chatWebSocket.onOpen(session2, "user2");
        assertEquals(2, ChatWebSocket.getOnlineUserCount());
        
        // 移除用户
        chatWebSocket.onClose("user1");
        assertEquals(1, ChatWebSocket.getOnlineUserCount());
    }

    @Test
    void testGetOnlineUsers() {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session2.getBasicRemote()).thenReturn(remoteEndpoint2);
        when(session1.isOpen()).thenReturn(true);
        when(session2.isOpen()).thenReturn(true);
        
        // 初始状态应该为空（由于@BeforeEach中的清理）
        Set<String> users = ChatWebSocket.getOnlineUsers();
        assertTrue(users.isEmpty());
        
        // 添加用户
        chatWebSocket.onOpen(session1, "user1");
        chatWebSocket.onOpen(session2, "user2");
        
        users = ChatWebSocket.getOnlineUsers();
        assertEquals(2, users.size());
        assertTrue(users.contains("user1"));
        assertTrue(users.contains("user2"));
    }

    @Test
    void testBroadcastMessage() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session2.getBasicRemote()).thenReturn(remoteEndpoint2);
        when(session1.isOpen()).thenReturn(true);
        when(session2.isOpen()).thenReturn(true);
        
        // 建立多个连接
        chatWebSocket.onOpen(session1, "user1");
        chatWebSocket.onOpen(session2, "user2");
        
        // 广播消息（排除user1）
        chatWebSocket.onMessage("Broadcast message", "user1");
        
        // 验证只有user2收到消息
        verify(remoteEndpoint2).sendText("用户 user1: Broadcast message");
        verify(remoteEndpoint1, never()).sendText(anyString());
    }

    @Test
    void testMessageServiceException() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session1.isOpen()).thenReturn(true);
        
        // 模拟MessageService抛出异常
        doThrow(new RuntimeException("Database error"))
            .when(messageService).saveMessage(anyString(), anyString(), anyString());
        
        // 建立连接
        chatWebSocket.onOpen(session1, "user1");
        
        // 发送消息不应该抛出异常
        assertDoesNotThrow(() -> {
            chatWebSocket.onMessage("Test message", "user1");
        });
    }

    @Test
    void testSendMessageException() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session1.isOpen()).thenReturn(true);
        
        // 建立连接
        chatWebSocket.onOpen(session1, "user1");
        
        // 模拟发送消息时抛出异常
        doThrow(new IOException("Send failed"))
            .when(remoteEndpoint1).sendText(anyString());
        
        // 发送消息应该返回false
        boolean result = ChatWebSocket.sendMessageToUser("user1", "message");
        
        assertFalse(result);
    }

    @Test
    void testMultipleUsersSameId() throws IOException {
        // 设置这个测试需要的mock行为
        when(session1.getBasicRemote()).thenReturn(remoteEndpoint1);
        when(session1.isOpen()).thenReturn(true);
        
        // 测试相同用户ID的多个连接（后连接覆盖前连接）
        Session session3 = mock(Session.class);
        RemoteEndpoint.Basic remoteEndpoint3 = mock(RemoteEndpoint.Basic.class);
        when(session3.getBasicRemote()).thenReturn(remoteEndpoint3);
        when(session3.isOpen()).thenReturn(true);
        
        // 建立第一个连接
        chatWebSocket.onOpen(session1, "user1");
        assertEquals(1, ChatWebSocket.getOnlineUserCount());
        
        // 建立第二个相同ID的连接（会覆盖第一个）
        chatWebSocket.onOpen(session3, "user1");
        assertEquals(1, ChatWebSocket.getOnlineUserCount()); // 数量不变，因为覆盖了
        
        // 发送消息应该发送给最新的连接
        boolean result = ChatWebSocket.sendMessageToUser("user1", "message");
        assertTrue(result);
        verify(remoteEndpoint3).sendText("message");
        verify(remoteEndpoint1, never()).sendText(anyString());
    }
}
